import numpy as np
from environment.GobangCheese import CheeseBoard
import random
import copy


UNUSED=0
def check_value(board):
    # 连续0 1 2 3 4 5个的数量分别是多少，初始化为0
    pos1_len={0:0,1:0,2:0,3:0,4:0,5:0}
    pos2_len={0:0,1:0,2:0,3:0,4:0,5:0}
    pos_dict=[None, pos1_len, pos2_len]
    
    for i in range(1, 16):
        pos_12=[UNUSED, 0, 0]
        for j in range(1, 16):    
            # 如果连续的1断了（黑棋断了）
            if board[i][j]!=1:
                pos_dict[1][min(pos_12[1], 5)]+=1
                if pos_12[1]>=5:
                    return pos_dict
                pos_12[1]=0
                
            # 连续的2断了
            if board[i][j]!=2:
                pos_dict[2][min(pos_12[2], 5)]+=1
                if pos_12[2]>=5:
                    return pos_dict
                pos_12[2]=0
                
            pos_12[board[i][j]]+=1
                

        for j in range(1,16):
            pos_12=[UNUSED, 0, 0]
            for i in range(1, 16):
                # 如果连续的1（黑棋断了）
                if board[i][j]!=1:
                    pos_dict[1][min(pos_12[1], 5)]+=1
                    if pos_12[1]>=5:
                        return pos_dict
                    pos_12[1]=0
                if board[i][j]!=2:
                    pos_dict[2][min(pos_12[2], 5)]+=1
                    if pos_12[2]>=5:
                        return pos_dict
                    pos_12[2]=0
                    
                pos_12[board[i][j]]+=1
                
        # 第一种倾斜方向
        for i in range(1, 16):
            for j in range(1, 16):
                pos_12=[UNUSED, 0, 0]
                for k in range(1, 16):
                    # 特殊情况剔除工作
                    if i+k>=16 or j+k>=16:
                        break
                    # 一般情况添加工作
                    pos_12[board[i+k][j+k]]+=1
                    
                    # 如果1断掉了
                    if board[i+k][j+k]!=1:
                        pos_dict[1][min(pos_12[1],5)]+=1
                        if pos_dict[1][5]>=5:
                            return pos_dict
                        pos_12[1]=0
                        
                    # 如果2断掉了
                    if board[i+k][j+k]!=2:
                        pos_dict[2][min(pos_12[2],5)]+=1
                        if pos_dict[2][5]>=5:
                            return pos_dict
                        pos_12[2]=0
                    
        for i in range(1, 16):
            for j in range(1, 16):
                pos_12=[UNUSED, 0, 0]
                for k in range(1, 16):
                    # 特殊情况剔除工作
                    if i+k>=16 or j+k>=16:
                        break
                    # 一般情况添加工作
                    pos_12[board[i+k][j-k]]+=1
                    
                    # 如果1断掉了
                    if board[i+k][j-k]!=1:
                        pos_dict[1][min(pos_12[1],5)]+=1
                        if pos_dict[1][5]>=5:
                            return pos_dict
                        pos_12[1]=0
                        
                    # 如果2断掉了
                    if board[i+k][j-k]!=2:
                        pos_dict[2][min(pos_12[2],5)]+=1
                        if pos_dict[2][5]>=5:
                            return pos_dict
                        pos_12[2]=0
    return pos_dict

def end_func(state):
    board=state.board
    values=check_value(board=board)
    return values[1][5]!=0 or values[2][5]!=0

def expanded_func(my_board):
    if np.any(my_board==0):
        return True
    return False

def naive_expand_func(old_state, hashval):
    # 基于五子棋规则，只可以在现有点的周边添加点
    state=copy.deepcopy(old_state)
    board=state.board
    player=state.player
    nonzero_indices = np.nonzero(board)
    nonzero_positions = list(zip(nonzero_indices[0], nonzero_indices[1]))
    '''
    锁定新节点的规则：
    dir_arr代表8个方向, 随机选择
    nonzero_positions, hashval决定了选择哪个节点
    综合方向和新节点，找到一个和某个当前节点相连接的新点
    '''
    dir_arr=[(0,1),(1,1),(1,0),(1,-1),(0,-1),(-1,-1),(-1,0),(-1,1)]
    choose_time=5
    #print(nonzero_positions)
    for i in range(choose_time):
        direction=random.randint(0,7)
        hashval=hashval%len(nonzero_positions)
        node=nonzero_positions[hashval]
        new_node=[node[0]+dir_arr[direction][0], node[1]+dir_arr[direction][1]]
        if new_node[0]>0 and new_node[1]>0 and new_node[0]<len(board) and new_node[0]<len(board):
            if board[new_node[0]][new_node[1]]==0:
                board[new_node[0]][new_node[1]]=player
                break
        hashval+=1
    return state    # 如果探索失败，返回原始状态。这样就会被MCTS的查重机制识别到

def naive_value_func(state):
    board=state.board
    value=check_value(board=board)
    # 接下来，计算value
    def value_func(value_dict):
        return value_dict[3]*1+value_dict[4]*3+value_dict[5]*6
    return value_func(value[2])-value_func(value[1])